import re
import math
from collections import defaultdict

# 数据预处理
def preprocess_text(text):
    text = re.sub(r'[^a-zA-Z\s]', '', text).lower()
    words = text.split()
    return words

# 硬编码的数据集
train_data = [
    "this is the first document",
    "this document is the second document",
    "and this is the third one",
    "is this the first document",
]

train_labels = [0, 1, 0, 0]  # 假设0和1是两个不同的类别

test_data = [
    "this is a test document",
    "this document will be classified",
    "another test document",
]

test_labels = [0, 1, 0]  # 相应的测试标签

# 训练朴素贝叶斯模型
def train_naive_bayes(data, labels):
    label_counts = defaultdict(int)
    word_counts = defaultdict(lambda: defaultdict(int))
    total_words = defaultdict(int)
    num_docs = len(data)
    
    for words, label in zip(data, labels):
        label_counts[label] += 1
        for word in words:
            word_counts[label][word] += 1
            total_words[label] += 1
    
    prior_probabilities = {label: math.log(count / num_docs) for label, count in label_counts.items()}
    conditional_probabilities = {}
    
    for label, word_count in word_counts.items():
        conditional_probabilities[label] = {}
        for word, count in word_count.items():
            conditional_probabilities[label][word] = math.log((count + 1) / (total_words[label] + len(word_counts[label])))
    
    return prior_probabilities, conditional_probabilities

# 预测函数
def predict(doc, prior_probabilities, conditional_probabilities):
    words = doc
    scores = {}
    
    for label, prior in prior_probabilities.items():
        score = prior
        for word in words:
            if word in conditional_probabilities[label]:
                score += conditional_probabilities[label][word]
            else:
                score += math.log(0.01)  # Add a small value for words not seen in training
        scores[label] = score
    
    return max(scores, key=scores.get)

# 计算分类报告
def classification_report(true_labels, predicted_labels, categories):
    correct_predictions = sum(1 for true, pred in zip(true_labels, predicted_labels) if true == pred)
    total_predictions = len(true_labels)
    accuracy = correct_predictions / total_predictions if total_predictions > 0 else 0
    
    print("Classification Report:")
    print(f"Accuracy: {accuracy:.2f}")

# 主函数
if __name__ == "__main__":
    # 加载数据集
    train_docs = [preprocess_text(doc) for doc in train_data]
    test_docs = [preprocess_text(doc) for doc in test_data]
    
    # 获取类别
    categories = sorted(set(train_labels))
    
    # 训练朴素贝叶斯模型
    prior_probabilities, conditional_probabilities = train_naive_bayes(train_docs, train_labels)
    
    # 预测测试集
    predictions = [predict(doc, prior_probabilities, conditional_probabilities) for doc in test_docs]
    
    # 打印分类报告
    classification_report(test_labels, predictions, categories)
